import torch
from torch import nn
from typing import Tuple


class Gru(nn.Module):
    def __init__(self, dims: Tuple[int, int, int], num_layer, dropout=0.1):
        super(Gru, self).__init__()
        self.input_size, self.hidden_size, self.output_size = dims
        self.num_layer = num_layer

        self.gru = nn.GRU(self.input_size, self.hidden_size, self.num_layer, batch_first=True, dropout=dropout)
        self.fc = nn.Linear(self.hidden_size, self.output_size)
        self.h0 = nn.Parameter(torch.zeros(self.num_layer, 1, self.hidden_size)).to('cuda')


    def forward(self, x):
        # 初始化隐藏状态
        h0 = torch.zeros(self.num_layer, x.size(0), self.hidden_size).to(x.device)

        # h0 = self.h0.repeat(1, x.shape[0], 1)
        # 通过GRU层传递输入和隐藏状态
        out, _ = self.gru(x, h0)
        # 这里改为把整个序列都返回
        out = self.fc(out)
        return out
